/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2022-2022. All rights reserved.
 * Description: 集合通信算子kernel
 * Author: lilianlin
 * Create: 2022-07-15
 */
#ifndef HCOM_LAUNCH_KERNEL_H
#define HCOM_LAUNCH_KERNEL_H
#include "hccl/base.h"
#include "hcom_build_graph.h"
#include "register/kernel_registry.h"
#include "register/kernel_registry_impl.h"
#include "exe_graph/lowering/shape_utils.h"
#include "exe_graph/runtime/tensor_data_utils.h"

namespace hccl {
constexpr int INPUT_INDEX_2 = 2;
constexpr int INPUT_INDEX_3 = 3;
constexpr int INPUT_INDEX_4 = 4;
constexpr int OUTPUT_INDEX_0 = 0;
constexpr int OUTPUT_INDEX_1 = 1;
constexpr int INDEX_PARA_2 = 2;
constexpr int INDEX_PARA_3 = 3;

constexpr u64 SHAPE_INFO_COUNT = 100; /* 动态send recv协商的shape info数据长度 */

struct HcomOpLaunchArgs {
    struct HcomOpAttr opAttr;
    void* stream;
    uint32_t inputNum;
    uint32_t outputNum;
    std::vector<void*> inputAddrs;
    std::vector<void*> outputAddrs;
    std::vector<gert::Shape> inputShapes;
    std::vector<gert::Shape> outputShapes;
};

ge::graphStatus LaunchHcomKernel(gert::KernelContext *context);
HcclResult GetHcomOpLaunchArgs(gert::KernelContext *context, HcomOpLaunchArgs& args);
HcclResult RegisterAddInfosFromLaunchArgs(HcomOpLaunchArgs& launchArgs, vector<HcomRemoteAccessAddrInfo>& addrInfos);
HcclResult ExtractRemoteReadAddrInfos(HcomOpLaunchArgs& launchArgs, vector<HcomRemoteAccessAddrInfo>& addrInfos);
HcclResult ExtractRemoteAddrInfosWithOffset(HcomOpLaunchArgs& launchArgs, vector<HcomRemoteAccessAddrInfo>& addrInfos);
HcclResult ExtractRemoteWriteAddsInfos(HcomOpLaunchArgs& launchArgs, vector<HcomRemoteAccessAddrInfo>& addrInfos);
HcclResult ExtractRemoteScatterWriteAddrInfos(HcomOpLaunchArgs& launchArgs,
    vector<HcomRemoteAccessAddrInfo>& addrInfos);
ge::graphStatus LaunchHcomKernelInitComm(gert::KernelContext *context);

ge::graphStatus BuildHcclOutputShapeOutputs(const ge::FastNode *node, gert::KernelContext *context);
HcclResult GetRecvOpLaunchArgs(gert::KernelContext *context, HcomOpLaunchArgs& args);
ge::graphStatus HcomGetRecvBeforeKernel(HcomOpLaunchArgs& args, std::vector<int64_t>& recvShape);
ge::graphStatus LaunchRecvKernel(gert::KernelContext *context);
ge::graphStatus LaunchRemotelookupKernel(gert::KernelContext *context);
HcclResult HcomReduceScatterVKernel(HcomOpLaunchArgs& launchArgs);
HcclResult HcomAllGatherVKernel(HcomOpLaunchArgs& launchArgs);
} // namespace hccl
#endif // HCOM_LAUNCH_KERNEL_H
